import re


class FunctionManager:
    """管理自定义函数的类，支持解析函数定义并存储可调用对象"""

    def __init__(self):
        self.functions = {}  # 存储函数名到 (参数列表, 表达式) 的映射

    def parse_function_definition(self, func_def: str):
        """
        解析函数定义字符串（如 "f(x)=2x"）

        参数:
        func_def (str): 函数定义字符串，格式要求：函数名(参数列表)=表达式
        """
        # 正则匹配函数定义格式：函数名(参数)=表达式
        pattern = r'^(\w+)\(([\w, ]+)\)\s*=\s*(.+?)(?:\s*:\s*(.+))?$'
        match = re.match(pattern, func_def)
        if not match:
            raise ValueError("函数定义格式错误，示例：f(x)=2x 或 g(a,b)=a+b")

        func_name, args_str, expr_part, conditions_str = match.groups()
        # 处理可能的None conditions_str
        conditions_str = conditions_str or ''

        # 解析表达式部分（移除可能的条件分隔符）
        expr = expr_part.strip()
        
        args = [arg.strip() for arg in args_str.split(',') if arg.strip()]

        # 验证参数合法性（仅允许字母）
        for arg in args:
            if not re.match(r'^\w+$', arg):
                raise ValueError(f"参数 '{arg}' 格式错误，仅允许字母数字下划线")

        param_conditions = {}
        return_condition = None

        if conditions_str:
            conditions = [cond.strip() for cond in
                          re.findall(r'[^,]*<<(?:[^()]*\([^()]*\)[^()]*)*[^,]*', conditions_str) if cond.strip()]
            for cond in conditions:
                if '<<' not in cond:
                    raise ValueError(f"条件 '{cond}' 格式错误，必须包含 '<<' 操作符")

                var_part, range_part = cond.split('<<', 1)
                var_part = var_part.strip()
                range_part = range_part.strip()

                # 检查是否为返回值条件 (格式为 函数名(参数列表))
                expected_return_var = f"{func_name}({', '.join(args)})"
                if var_part == expected_return_var:
                    if return_condition is not None:
                        raise ValueError("只能定义一个返回值条件")
                    return_condition = range_part
                elif var_part in args:
                    if var_part in param_conditions:
                        raise ValueError(f"参数 '{var_part}' 不能有多个条件")
                    param_conditions[var_part] = range_part
                else:
                    raise ValueError(f"条件 '{cond}' 中的变量 '{var_part}' 不是函数参数或返回值")

        self.functions[func_name] = (args, expr, param_conditions, return_condition)

    def evaluate_function(self, func_name: str, *args_values):
        """
        调用自定义函数（支持递归调用）

        参数:
        func_name (str): 函数名
        args_values (tuple): 传入的参数值

        返回:
        计算结果（可能是 CountModel 实例）
        """
        if func_name not in self.functions:
            raise ValueError(f"函数 '{func_name}' 未定义")

        args, expr, param_conditions, return_condition = self.functions[func_name]
        if len(args) != len(args_values):
            raise ValueError(f"函数 '{func_name}' 需要 {len(args)} 个参数，实际传入 {len(args_values)} 个")

        # 创建参数替换字典
        substitutions = {arg: value for arg, value in zip(args, args_values)}
        # 创建命名空间，包含参数和函数自身引用以支持递归
        namespace = substitutions.copy()
        namespace[func_name] = lambda *args: self.evaluate_function(func_name, *args)

        # 检查参数是否符合作用域条件
        for arg, value in zip(args, args_values):
            if arg in param_conditions:
                range_str = param_conditions[arg]
                if not self._check_range(value, range_str, namespace):
                    raise ValueError(f"参数 '{arg}' 的值 {value} 超出范围 {range_str}")

        try:
            # 使用eval计算表达式，限制命名空间确保安全
            result = eval(expr, {"__builtins__": None}, namespace)
            result = complex(result)
            # 检查返回值是否符合作用域条件
            if return_condition:
                if not self._check_range(result, return_condition, namespace):
                    raise ValueError(f"函数 '{func_name}' 的返回值 {result} 超出范围 {return_condition}")
            return result
        except Exception as e:
            raise ValueError(f"函数计算错误: {e}")

    def _check_range(self, value, range_str, namespace):
        """检查值是否在指定范围内"""
        # 支持的范围格式: (a,b), [a,b], (a,b], [a,b)
        pattern = r'^([\\(\\[])\s*([^,]+?)\s*,\s*([^,]+?)\s*([\\)\\]])$'
        match = re.match(pattern, range_str)
        if not match:
            raise ValueError(f"无效的范围格式: {range_str}，预期格式如 (a,b), [a,b], (a,b] 或 [a,b)")

        left_bracket, lower_str, upper_str, right_bracket = match.groups()

        # 评估上下限表达式（使用当前命名空间）
        try:
            lower = eval(lower_str, {"__builtins__": None}, namespace)
            upper = eval(upper_str, {"__builtins__": None}, namespace)
        except Exception as e:
            raise ValueError(f"评估范围边界时出错: {e}")

        # 检查左边界
        left_inclusive = (left_bracket == '[')
        if left_inclusive:
            left_ok = (value >= lower)
        else:
            left_ok = (value > lower)

        # 检查右边界
        right_inclusive = (right_bracket == ']')
        if right_inclusive:
            right_ok = (value <= upper)
        else:
            right_ok = (value < upper)

        return left_ok and right_ok
